# !/usr/bin/env python3
# coding=utf8
"""
sqlite 里的 dbbardata 数据表的载入器(从 dbbardata 把数据查出来, 让其他程序使用的类)
"""
import datetime
import pandas
import pathlib
import sqlite3
from typing import Dict, List, Set, Tuple, Optional, Union, Callable

from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.utility import get_file_path


class DbBarDataLoader(object):
    """
    sqlite 里的数据表 dbbardata 的载入器
    """

    DBPATH: pathlib.Path = get_file_path("database.db")

    def __init__(
        self,
        exchange: Exchange,
        symbol: str,
        interval: Interval,
        start: datetime.datetime,
        end: datetime.datetime,
        preload_num: int,
    ) -> None:
        """"""
        self.exchange: Exchange = exchange  # 交易所
        self.symbol: str = symbol  # 代码
        self.interval: Interval = interval  # K线周期
        self.start: datetime.datetime = start  # 开始日期
        self.end: datetime.datetime = end  # 结束日期
        self.prelod_num: int = preload_num  # 预加载多少条数据

        self.data: pandas.DataFrame = None  # 数据

    def execute(self) -> pandas.DataFrame:
        """"""
        self.data: pandas.DataFrame = self.load_data(
            exchange=self.exchange,
            symbol=self.symbol,
            interval=self.interval,
            start=self.start,
            end=self.end,
            preload_num=self.prelod_num,
        )
        return self.data

    @classmethod
    def preload_query(
        cls,
        exchange: Exchange,
        symbol: str,
        interval: Interval,
        start: datetime.datetime,
        end: datetime.datetime,
        num: int,
    ) -> Tuple[datetime.datetime, int]:
        """
        我想加载时间点 start 前面的 num 条数据,
        我想查询 start 前面的数据量够不够 num 条, 以及第 num 条数据的时间戳是什么,
        """
        if num == 0:
            return start, num

        sql_statement: str = f"""
            SELECT MIN(datetime) AS min_datetime, COUNT(datetime) AS cnt_datetime FROM ( 
                SELECT datetime FROM dbbardata WHERE 
                    exchange = '{exchange.value}' AND symbol = '{symbol}' 
                    AND interval = '{interval.value}' 
                    AND datetime < '{start.strftime('%Y-%m-%d %H:%M:%S')}' 
                    ORDER BY datetime DESC LIMIT '{num}'
            );"""

        with sqlite3.connect(database=cls.DBPATH) as connection:
            cursor = connection.cursor()
            cursor.execute(sql_statement)

            head = [col[0] for col in cursor.description]
            results = [dict(zip(head, line)) for line in cursor.fetchall()]

            assert len(results) == 1
            result: dict = results[0]

            min_datetime: str = result["min_datetime"]
            cnt_datetime: int = result["cnt_datetime"]

            min_datetime: datetime.datetime = datetime.datetime.strptime(min_datetime, "%Y-%m-%d %H:%M:%S")

            return min_datetime, cnt_datetime

    @classmethod
    def load_data(
        cls,
        exchange: Exchange,
        symbol: str,
        interval: Interval,
        start: datetime.datetime,
        end: datetime.datetime,
        preload_num: int,
    ) -> pandas.DataFrame:
        """"""
        preload_start, preload_cnt = cls.preload_query(
            exchange=exchange,
            symbol=symbol,
            interval=interval,
            start=start,
            end=end,
            num=preload_num,
        )
        if preload_cnt != preload_num:
            raise RuntimeError(f"预加载数据量不足, 期望预加载{preload_num}条, 实际只能预加载{preload_cnt}条,")

        sql_statement: str = f"""
            SELECT * FROM dbbardata WHERE 
                exchange = '{exchange.value}' AND symbol = '{symbol}' 
                AND interval = '{interval.value}' 
                AND datetime >= '{preload_start.strftime('%Y-%m-%d %H:%M:%S')}' 
                AND datetime <= '{end.strftime('%Y-%m-%d %H:%M:%S')}';"""

        with sqlite3.connect(database=cls.DBPATH) as connection:
            df: pandas.DataFrame = pandas.read_sql(
                sql=sql_statement,
                con=connection,
            )

            # series_timestamp = pandas.to_datetime(df["datetime"])
            # assert isinstance(series_timestamp, pandas.Series)
            # assert isinstance(series_timestamp[0], pandas.Timestamp)

            df["datetime"] = pandas.to_datetime(df["datetime"])  # 字符串转时间戳

            # inplace: 是否在原 DataFrame 上修改, 默认为 False
            # verify_integrity: 是否检查索引有无重复, 默认为 False, 若设置为 True 会影响程序性能,
            df.set_index("datetime", inplace=True, verify_integrity=False)

            return df


if __name__ == "__main__":
    loader: DbBarDataLoader = DbBarDataLoader(
        exchange=Exchange.SHFE,
        symbol="AUL8",
        interval=Interval.MINUTE,
        start=datetime.datetime(2022, 12, 1),
        end=datetime.datetime(2022, 12, 1),
        preload_num=100,
    )
    loader.execute()
